{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "%matplotlib inline"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "\n=================================\nPlot graphs' barycenter using FGW\n=================================\n\nThis example illustrates the computation barycenter of labeled graphs using FGW\n\nRequires networkx >=2\n\n.. [18] Vayer Titouan, Chapel Laetitia, Flamary R{'e}mi, Tavenard Romain\n      and Courty Nicolas\n    \"Optimal Transport for structured data with application on graphs\"\n    International Conference on Machine Learning (ICML). 2019.\n\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "# Author: Titouan Vayer <titouan.vayer@irisa.fr>\n#\n# License: MIT License\n\n#%% load libraries\nimport numpy as np\nimport matplotlib.pyplot as plt\nimport networkx as nx\nimport math\nfrom scipy.sparse.csgraph import shortest_path\nimport matplotlib.colors as mcol\nfrom matplotlib import cm\nfrom ot.gromov import fgw_barycenters\n#%% Graph functions\n\n\ndef find_thresh(C, inf=0.5, sup=3, step=10):\n    \"\"\" Trick to find the adequate thresholds from where value of the C matrix are considered close enough to say that nodes are connected\n        Tthe threshold is found by a linesearch between values \"inf\" and \"sup\" with \"step\" thresholds tested.\n        The optimal threshold is the one which minimizes the reconstruction error between the shortest_path matrix coming from the thresholded adjency matrix\n        and the original matrix.\n    Parameters\n    ----------\n    C : ndarray, shape (n_nodes,n_nodes)\n            The structure matrix to threshold\n    inf : float\n          The beginning of the linesearch\n    sup : float\n          The end of the linesearch\n    step : integer\n            Number of thresholds tested\n    \"\"\"\n    dist = []\n    search = np.linspace(inf, sup, step)\n    for thresh in search:\n        Cprime = sp_to_adjency(C, 0, thresh)\n        SC = shortest_path(Cprime, method='D')\n        SC[SC == float('inf')] = 100\n        dist.append(np.linalg.norm(SC - C))\n    return search[np.argmin(dist)], dist\n\n\ndef sp_to_adjency(C, threshinf=0.2, threshsup=1.8):\n    \"\"\" Thresholds the structure matrix in order to compute an adjency matrix.\n    All values between threshinf and threshsup are considered representing connected nodes and set to 1. Else are set to 0\n    Parameters\n    ----------\n    C : ndarray, shape (n_nodes,n_nodes)\n        The structure matrix to threshold\n    threshinf : float\n        The minimum value of distance from which the new value is set to 1\n    threshsup : float\n        The maximum value of distance from which the new value is set to 1\n    Returns\n    -------\n    C : ndarray, shape (n_nodes,n_nodes)\n        The threshold matrix. Each element is in {0,1}\n    \"\"\"\n    H = np.zeros_like(C)\n    np.fill_diagonal(H, np.diagonal(C))\n    C = C - H\n    C = np.minimum(np.maximum(C, threshinf), threshsup)\n    C[C == threshsup] = 0\n    C[C != 0] = 1\n\n    return C\n\n\ndef build_noisy_circular_graph(N=20, mu=0, sigma=0.3, with_noise=False, structure_noise=False, p=None):\n    \"\"\" Create a noisy circular graph\n    \"\"\"\n    g = nx.Graph()\n    g.add_nodes_from(list(range(N)))\n    for i in range(N):\n        noise = float(np.random.normal(mu, sigma, 1))\n        if with_noise:\n            g.add_node(i, attr_name=math.sin((2 * i * math.pi / N)) + noise)\n        else:\n            g.add_node(i, attr_name=math.sin(2 * i * math.pi / N))\n        g.add_edge(i, i + 1)\n        if structure_noise:\n            randomint = np.random.randint(0, p)\n            if randomint == 0:\n                if i <= N - 3:\n                    g.add_edge(i, i + 2)\n                if i == N - 2:\n                    g.add_edge(i, 0)\n                if i == N - 1:\n                    g.add_edge(i, 1)\n    g.add_edge(N, 0)\n    noise = float(np.random.normal(mu, sigma, 1))\n    if with_noise:\n        g.add_node(N, attr_name=math.sin((2 * N * math.pi / N)) + noise)\n    else:\n        g.add_node(N, attr_name=math.sin(2 * N * math.pi / N))\n    return g\n\n\ndef graph_colors(nx_graph, vmin=0, vmax=7):\n    cnorm = mcol.Normalize(vmin=vmin, vmax=vmax)\n    cpick = cm.ScalarMappable(norm=cnorm, cmap='viridis')\n    cpick.set_array([])\n    val_map = {}\n    for k, v in nx.get_node_attributes(nx_graph, 'attr_name').items():\n        val_map[k] = cpick.to_rgba(v)\n    colors = []\n    for node in nx_graph.nodes():\n        colors.append(val_map[node])\n    return colors"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Generate data\n-------------\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "#%% circular dataset\n# We build a dataset of noisy circular graphs.\n# Noise is added on the structures by random connections and on the features by gaussian noise.\n\n\nnp.random.seed(30)\nX0 = []\nfor k in range(9):\n    X0.append(build_noisy_circular_graph(np.random.randint(15, 25), with_noise=True, structure_noise=True, p=3))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Plot data\n---------\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "#%% Plot graphs\n\nplt.figure(figsize=(8, 10))\nfor i in range(len(X0)):\n    plt.subplot(3, 3, i + 1)\n    g = X0[i]\n    pos = nx.kamada_kawai_layout(g)\n    nx.draw(g, pos=pos, node_color=graph_colors(g, vmin=-1, vmax=1), with_labels=False, node_size=100)\nplt.suptitle('Dataset of noisy graphs. Color indicates the label', fontsize=20)\nplt.show()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Barycenter computation\n----------------------\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "#%% We compute the barycenter using FGW. Structure matrices are computed using the shortest_path distance in the graph\n# Features distances are the euclidean distances\nCs = [shortest_path(nx.adjacency_matrix(x)) for x in X0]\nps = [np.ones(len(x.nodes())) / len(x.nodes()) for x in X0]\nYs = [np.array([v for (k, v) in nx.get_node_attributes(x, 'attr_name').items()]).reshape(-1, 1) for x in X0]\nlambdas = np.array([np.ones(len(Ys)) / len(Ys)]).ravel()\nsizebary = 15  # we choose a barycenter with 15 nodes\n\nA, C, log = fgw_barycenters(sizebary, Ys, Cs, ps, lambdas, alpha=0.95, log=True)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Plot Barycenter\n-------------------------\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "#%% Create the barycenter\nbary = nx.from_numpy_matrix(sp_to_adjency(C, threshinf=0, threshsup=find_thresh(C, sup=100, step=100)[0]))\nfor i, v in enumerate(A.ravel()):\n    bary.add_node(i, attr_name=v)\n\n#%%\npos = nx.kamada_kawai_layout(bary)\nnx.draw(bary, pos=pos, node_color=graph_colors(bary, vmin=-1, vmax=1), with_labels=False)\nplt.suptitle('Barycenter', fontsize=20)\nplt.show()"
      ]
    }
  ],
  "metadata": {
    "kernelspec": {
      "display_name": "Python 3",
      "language": "python",
      "name": "python3"
    },
    "language_info": {
      "codemirror_mode": {
        "name": "ipython",
        "version": 3
      },
      "file_extension": ".py",
      "mimetype": "text/x-python",
      "name": "python",
      "nbconvert_exporter": "python",
      "pygments_lexer": "ipython3",
      "version": "3.6.8"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}